package com.example.tls_mitm_learn;

import android.annotation.SuppressLint;
import android.net.VpnService;

import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509ExtensionUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;

import javax.net.ssl.*;
import java.io.*;
import java.math.BigInteger;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.*;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

public class Socks5Proxy {

    // ==========================================
    // Constants
    // ==========================================
    public static final int LISTEN_PORT = 9090;
    private static final int TLS_PORT = 443;

    private static final ExecutorService THREAD_POOL = Executors.newCachedThreadPool();

    private static final String CA_P12_FILE = "myproxy-ca.p12";
    public static final String CERT_FILE = "ca_cert.crt";

    private static final char[] CA_PASSWORD = "password".toCharArray();
    private static final char[] KEY_PASSWORD = "password".toCharArray();
    //证书
    private static volatile X509Certificate caCertificate;
    //私钥
    private static volatile PrivateKey caPrivateKey;

    private static final Map<String, KeyStore> DOMAIN_KEYSTORE_CACHE = new ConcurrentHashMap<>();

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    // ==========================================
    // Main
    // ==========================================
    public static void main(String[] args) {
        try (ServerSocket serverSocket = new ServerSocket(LISTEN_PORT)) {
            System.out.println("Listening on port " + LISTEN_PORT);
            //初始化或者加载CA证书
            initializeOrLoadCA();
            System.out.println("Loaded CA: " + caCertificate.getSubjectX500Principal());

            while (!Thread.interrupted()) {
                Socket client = serverSocket.accept();
                THREAD_POOL.execute(() -> handleClient(client));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    @SuppressLint("StaticFieldLeak")
    static VpnService _vpnService;
    public static void start(VpnService vpnService) {
        _vpnService = vpnService;
        try (ServerSocket serverSocket = new ServerSocket(LISTEN_PORT)) {
            System.out.println("Listening on port " + LISTEN_PORT);
            //初始化或者加载CA证书
            initializeOrLoadCA();
            System.out.println("Loaded CA: " + caCertificate.getSubjectX500Principal());

            while (!Thread.interrupted()) {
                Socket client = serverSocket.accept();
                THREAD_POOL.execute(() -> handleClient(client));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    // ==========================================
    // Client Handler
    // ==========================================
    private static void handleClient(Socket client) {
        try (client;
             InputStream in = client.getInputStream();
             OutputStream out = client.getOutputStream()) {

            int version = in.read();
            if (version != 0x05) return;

            int nMethods = in.read();
            readNBytesCompat(in,nMethods);

            out.write(new byte[]{0x05, 0x00}); // NO AUTH
            out.flush();

            int ver = in.read();
            int cmd = in.read();
            in.read(); // RSV
            int addrType = in.read();
            String host = readHost(in, addrType);
            int port = ((in.read() & 0xFF) << 8) | (in.read() & 0xFF);

            if (cmd != 0x01) throw new IOException("Only CONNECT command is supported");

            out.write(new byte[]{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0});
            out.flush();
            MyVpnService myVpnService = (MyVpnService) _vpnService;
            host = myVpnService.ipAndHostMap.get(host);
            _vpnService.protect(client);
            if (port == TLS_PORT) {//443端口 https
                handleTlsMitm(client, host, port);
            } else { //其他端口 直接转发
                tcpPipe(client, host, port);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static String readHost(InputStream in, int addrType) throws IOException {
        if (addrType == 0x01) { // IPv4
            byte[] ip = readNBytesCompat(in,4);
            return (ip[0] & 0xFF) + "." + (ip[1] & 0xFF) + "." + (ip[2] & 0xFF) + "." + (ip[3] & 0xFF);
        } else if (addrType == 0x03) { // Domain
            int len = in.read();
            byte[] name = readNBytesCompat(in,len);
            return new String(name);
        } else {
            throw new IOException("Unsupported address type: " + addrType);
        }
    }

    // ==========================================
    // TCP & TLS Forwarding
    // ==========================================
    private static void tcpPipe(Socket client, String host, int port) throws Exception {

        try (Socket server = new Socket();client) {
            server.bind(null);
            _vpnService.protect(server);
            server.connect(new InetSocketAddress(host, port));

            ExecutorCompletionService<Void> ecs = new ExecutorCompletionService<>(THREAD_POOL);
            ecs.submit(() -> { transferData(client, server); return null; });
            ecs.submit(() -> { transferData(server, client); return null; });
            // 等待任意一个任务完成
            ecs.take();
        }
    }

    private static void handleTlsMitm(Socket client, String host, int port) throws Exception {
        //通过CA证书为目标域名生成证书
        KeyStore ks = getOrCreateDomainKeyStore(host);

        SSLContext serverCtx = SSLContext.getInstance("TLS");
        KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        kmf.init(ks, KEY_PASSWORD);
        serverCtx.init(kmf.getKeyManagers(), null, null);
        //向客户端发起TLS握手
        SSLSocket sslServer = (SSLSocket) serverCtx.getSocketFactory().createSocket(client,
                client.getInetAddress().getHostAddress(), client.getPort(), true);
        _vpnService.protect(sslServer);
        sslServer.setUseClientMode(false);
        sslServer.setNeedClientAuth(false);
        sslServer.startHandshake();

        //发送https
        SSLSocket sslClient = (SSLSocket) SSLSocketFactory.getDefault().createSocket();
        sslClient.bind(null);
        _vpnService.protect(sslClient);
        sslClient.connect(new InetSocketAddress(host, port));
        sslClient.setUseClientMode(true);
        sslClient.startHandshake();

        //双向转发数据
        try (sslServer; sslClient) {
            ExecutorCompletionService<Void> ecs = new ExecutorCompletionService<>(THREAD_POOL);
            ecs.submit(() -> { transferData(sslServer, sslClient); return null; });
            ecs.submit(() -> { transferData(sslClient, sslServer); return null; });

            // 等待任意一个任务完成
            ecs.take();
        }
    }

    /**
     * 通用版 readNBytes，兼容 Android 全版本。
     * 与 Java 9+ 的 InputStream#readNBytes 行为一致。
     */
    private static byte[] readNBytesCompat(InputStream in, int len) throws IOException {
        if (len < 0) throw new IllegalArgumentException("len must be non-negative");
        byte[] buffer = new byte[len];
        int total = 0;
        while (total < len) {
            int read = in.read(buffer, total, len - total);
            if (read == -1) break; // EOF
            total += read;
        }
        if (total == len) return buffer;
        // 如果实际读取不足，则返回已读部分（与 readNBytes 相同语义）
        byte[] result = new byte[total];
        System.arraycopy(buffer, 0, result, 0, total);
        return result;
    }

    private static void transferData(Socket src, Socket dst) {
        byte[] buffer = new byte[8192];
        int bytesRead;
        try (InputStream in = src.getInputStream(); OutputStream out = dst.getOutputStream()) {
            while ((bytesRead = in.read(buffer)) != -1) {
                out.write(buffer, 0, bytesRead);
                out.flush();
                String s = new String(buffer, 0, bytesRead);
                System.out.println("[Forward] " + s.replaceAll("\r\n", "\\\\r\\n"));
            }
        } catch (IOException ignored) {}
    }

    // ==========================================
    // Certificate Management
    // ==========================================
    private static KeyStore getOrCreateDomainKeyStore(String hostname) {
        return DOMAIN_KEYSTORE_CACHE.computeIfAbsent(hostname, h -> {
            try {
                KeyPair leafKeyPair = generateRSAKeyPair(2048);
                X509Certificate leafCert = generateServerCertificate(h, leafKeyPair, caCertificate, caPrivateKey);

                KeyStore ks = KeyStore.getInstance("PKCS12");
                ks.load(null, null);
                ks.setKeyEntry("alias", leafKeyPair.getPrivate(), KEY_PASSWORD,
                        new java.security.cert.Certificate[]{leafCert, caCertificate});
                return ks;
            } catch (Exception e) {
                throw new RuntimeException("Failed to generate certificate for domain " + h, e);
            }
        });
    }

    private static KeyPair generateRSAKeyPair(int keySize) throws Exception {
        KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
        generator.initialize(keySize);
        return generator.generateKeyPair();
    }

    private static X509Certificate generateServerCertificate(String domain, KeyPair keyPair,
                                                             X509Certificate caCert, PrivateKey caKey) throws Exception {
        long now = System.currentTimeMillis();
        Date notBefore = new Date(now - 24L * 3600 * 1000);
        Date notAfter = new Date(now + 365L * 24 * 3600 * 1000);
        //证书信息
        X500Name issuer = new X500Name(caCert.getSubjectX500Principal().getName());
        X500Name subject = new X500Name("CN=" + domain);
        BigInteger serial = new BigInteger(64, new SecureRandom()).abs();
        SubjectPublicKeyInfo spki = SubjectPublicKeyInfo.getInstance(keyPair.getPublic().getEncoded());

        X509v3CertificateBuilder builder = new X509v3CertificateBuilder(
                issuer, serial, notBefore, notAfter, subject, spki);
        //添加扩展
        builder.addExtension(Extension.basicConstraints, true, new BasicConstraints(false));//非CA
        builder.addExtension(Extension.keyUsage, true,
                new KeyUsage(KeyUsage.digitalSignature | KeyUsage.keyEncipherment));
        builder.addExtension(Extension.extendedKeyUsage, true,
                new ExtendedKeyUsage(KeyPurposeId.id_kp_serverAuth));
        builder.addExtension(Extension.subjectAlternativeName, false,
                new GeneralNames(new GeneralName(GeneralName.dNSName, domain)));//域名

        JcaX509ExtensionUtils extUtils = new JcaX509ExtensionUtils();
        builder.addExtension(Extension.subjectKeyIdentifier, false, extUtils.createSubjectKeyIdentifier(keyPair.getPublic()));
        builder.addExtension(Extension.authorityKeyIdentifier, false, extUtils.createAuthorityKeyIdentifier(caCert.getPublicKey()));
        //签名
        ContentSigner signer = new JcaContentSignerBuilder("SHA256withRSA").build(caKey);
        X509CertificateHolder holder = builder.build(signer);
        return new JcaX509CertificateConverter().getCertificate(holder);
    }

    private static X509Certificate generateRootCA(String cn, KeyPair keyPair) throws Exception {
        long now = System.currentTimeMillis();
        Date notBefore = new Date(now - 24L * 3600 * 1000);
        Date notAfter = new Date(now + 3650L * 24 * 3600 * 1000);
        //证书信息
        X500Name subject = new X500Name("CN=" + cn);
        X500Name issuer = subject;
        BigInteger serial = new BigInteger(64, new SecureRandom()).abs();
        SubjectPublicKeyInfo spki = SubjectPublicKeyInfo.getInstance(keyPair.getPublic().getEncoded());

        X509v3CertificateBuilder builder = new X509v3CertificateBuilder(
                issuer, serial, notBefore, notAfter, subject, spki);
        //添加扩展
        builder.addExtension(Extension.basicConstraints, true, new BasicConstraints(true));
        builder.addExtension(Extension.keyUsage, true,
                new KeyUsage(KeyUsage.keyCertSign | KeyUsage.cRLSign));

        JcaX509ExtensionUtils extUtils = new JcaX509ExtensionUtils();
        builder.addExtension(Extension.subjectKeyIdentifier, false, extUtils.createSubjectKeyIdentifier(keyPair.getPublic()));
        builder.addExtension(Extension.authorityKeyIdentifier, false, extUtils.createAuthorityKeyIdentifier(keyPair.getPublic()));
        //签名
        ContentSigner signer = new JcaContentSignerBuilder("SHA256withRSA").build(keyPair.getPrivate());
        X509CertificateHolder holder = builder.build(signer);
        return new JcaX509CertificateConverter().getCertificate(holder);
    }

    private static void initializeOrLoadCA() throws Exception {
        File p12File = new File(_vpnService.getFilesDir(),CA_P12_FILE);
        if (p12File.exists()) {
            System.out.println("[*] Loading existing CA...");
            loadCAPkcs12();
            System.out.println("[+] CA loaded: " + caCertificate.getSubjectX500Principal());
        } else {
            System.out.println("[*] Generating new CA...");
            KeyPair kp = generateRSAKeyPair(2048);
            X509Certificate cert = generateRootCA("MyProxyCA", kp);
            saveCAPkcs12(kp.getPrivate(), cert);
            saveCACrt(cert);
            caCertificate = cert;
            caPrivateKey = kp.getPrivate();
            System.out.println("[+] CA generated and saved to: " + CA_P12_FILE);
            System.out.println("[!] Please import the CA certificate into OS/browser trust store.");
        }
    }

    private static void saveCAPkcs12(PrivateKey caKey, X509Certificate caCert) throws Exception {
        KeyStore ks = KeyStore.getInstance("PKCS12");
        ks.load(null, null);
        ks.setKeyEntry("ca", caKey, CA_PASSWORD, new java.security.cert.Certificate[]{caCert});
        try (FileOutputStream fos = new FileOutputStream(new File(_vpnService.getFilesDir(),CA_P12_FILE))) {
            ks.store(fos, CA_PASSWORD);
        }
    }

    // 保存只包含公钥的 crt 文件（给 Android 系统导入为 CA 证书）
    private static void saveCACrt(X509Certificate caCert) throws Exception {
        File crtFile = new File(_vpnService.getFilesDir(), CERT_FILE);
        try (FileOutputStream fos = new FileOutputStream(crtFile)) {
            fos.write(caCert.getEncoded());  // DER 编码
        }
    }

    private static void loadCAPkcs12() throws Exception {
        KeyStore ks = KeyStore.getInstance("PKCS12");
        try (FileInputStream fis = new FileInputStream(new File(_vpnService.getFilesDir(),CA_P12_FILE))) {
            ks.load(fis, CA_PASSWORD);
        }

        Key key = ks.getKey("ca", CA_PASSWORD);
        if (!(key instanceof PrivateKey)) throw new IllegalStateException("CA entry is not a private key");
        caPrivateKey = (PrivateKey) key;

        java.security.cert.Certificate cert = ks.getCertificate("ca");
        if (!(cert instanceof X509Certificate)) throw new IllegalStateException("CA certificate not X509");
        caCertificate = (X509Certificate) cert;
    }
}
